# encoding: utf-8

import numpy as np
from scipy import signal as Scisig
from scipy import signal
from scipy.signal import filtfilt, butter, sosfiltfilt

from utils import lfp_bandpass


# Author     : forward4<liang.ma@siat.ac.cn>
# Datetime   :
# Project    : monkey_neuralDecoding
# File       : dataProcessor.py
# description: 文件说明

def valid_data(func):
    def inner(*args, **kwargs):
        """
        :param args: arg[0] is an instance of DataProcessor; arg[1] is the signal to be processed.
        :param kwargs:
        :return:
        """
        verbose = kwargs.get('verbose', False)
        if verbose:
            print("validating data")

        _sig = args[1]
        assert isinstance(_sig, np.ndarray)

        if verbose:
            print("validating over.")
        return func(*args, **kwargs)

    return inner


class DataProcessor:
    def __init__(self):
        pass

    @valid_data
    def bandpass_filter(self, signal, lowcut=100, highcut=300, fs=30000):
        _rs = lfp_bandpass(signal, low=lowcut, high=highcut, fs=fs, order=2)
        return _rs

    @valid_data
    def lowpass_filter(self, signal, highcut=250, fs=30000):
        sos = butter(6, highcut, btype='lp', fs=fs, output='sos')
        filtered = sosfiltfilt(sos, signal)
        return filtered

    @valid_data
    def highpass_filter(self, signal, lowcut=300, fs=30000):
        pass

    @valid_data
    def notch_filter(self, signal, notched_fs=50, fs=30000, Q=0.2):
        """
        :param signal:
        :param notched_fs:
        :param fs:
        :param Q: quality factor.
        :return:
        """
        w0 = 2 * notched_fs / fs
        b, a = Scisig.iirnotch(notched_fs, Q=Q, fs=fs)
        return np.array(filtfilt(b, a, signal))

    @valid_data
    def get_statistics(self, signal, **kwargs):
        pass

    @valid_data
    def down_sampling(self, signal_data, factor):
        num_channels = signal_data.shape[0]
        resampled_data = np.zeros((num_channels, len(signal_data[0]) // factor))
        for i in range(num_channels):
            resampled_data[i] = signal.resample(signal_data[i], len(signal_data[i]) // factor)
        return resampled_data

    @valid_data
    def resampling(self,signal_data, length):
        num_channels = signal_data.shape[0]
        resampled_data = np.zeros((num_channels,length))
        for i in range(num_channels):
            resampled_data[i] = signal.resample(signal_data[i],length)
        return resampled_data


if __name__ == '__main__':
    data_processor = DataProcessor()
    sig = np.arange(20).reshape(4, 5)
    data_processor.notch_filter(sig)
    pass
